/*
 * wifi_connect.c
 *
 *  Created on: Mar 16, 2021
 *      Author: kychu
 */

#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "nvs_flash.h"
#include <unistd.h>
#include <sys/unistd.h>
#include "lwip/err.h"
#include "lwip/sys.h"

#include "cJSON.h"
#include "mdns.h"
#include "wifi_connect.h"
#include "lwip/apps/netbiosns.h"

#define WIFI_SSID_PASSWD_MAX_LEN 32

typedef struct {
  char host_name[WIFI_SSID_PASSWD_MAX_LEN];
  char sta_ssid[WIFI_SSID_PASSWD_MAX_LEN];
  char sta_passwd[WIFI_SSID_PASSWD_MAX_LEN];
  char ap_ssid[WIFI_SSID_PASSWD_MAX_LEN];
  char ap_passwd[WIFI_SSID_PASSWD_MAX_LEN];
} sys_wifi_cfg_t;

static const char *wifi_cfg_json_name[] = {
  "host_name",
  "sta_ssid",
  "sta_passwd",
  "ap_ssid",
  "ap_passwd"
};

static sys_wifi_cfg_t *cfg = NULL;
static char ipv4_ip_str[16] = "0.0.0.0";

#define EXAMPLE_ESP_MAXIMUM_RETRY  5
#define EXAMPLE_ESP_WIFI_CHANNEL   1
#define EXAMPLE_MAX_STA_CONN       4

/* FreeRTOS event group to signal when we are connected*/
static EventGroupHandle_t s_wifi_event_group;

/* The event group allows multiple bits for each event, but we only care about two events:
 * - we are connected to the AP with an IP
 * - we failed to connect after the maximum amount of retries */
#define WIFI_CONNECTED_BIT BIT0
#define WIFI_FAIL_BIT      BIT1

static const char *TAG = "wifi_conn";

static int s_retry_num = 0;

static void event_handler(void* arg, esp_event_base_t event_base,
                                int32_t event_id, void* event_data)
{
//  ESP_LOGI(TAG, "new event: %s, %d", event_base, event_id);
  if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START) {
    esp_wifi_connect();
  } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
    if (s_retry_num < EXAMPLE_ESP_MAXIMUM_RETRY) {
      esp_wifi_connect();
      s_retry_num ++;
      ESP_LOGI(TAG, "retry to connect to the AP");
    } else {
      xEventGroupSetBits(s_wifi_event_group, WIFI_FAIL_BIT);
    }
    ESP_LOGI(TAG,"connect to the AP fail");
  } else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
    ip_event_got_ip_t* event = (ip_event_got_ip_t*) event_data;
    memset(ipv4_ip_str, 0, 16);
    snprintf(ipv4_ip_str, 16, IPSTR, IP2STR(&event->ip_info.ip));
    ESP_LOGI(TAG, "got ip:%s", ipv4_ip_str);
    s_retry_num = 0;
    xEventGroupSetBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
  } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_AP_STACONNECTED) {
    wifi_event_ap_staconnected_t* event = (wifi_event_ap_staconnected_t*) event_data;
    ESP_LOGI(TAG, "station "MACSTR" join, AID=%d", MAC2STR(event->mac), event->aid);
  } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_AP_STADISCONNECTED) {
    wifi_event_ap_stadisconnected_t* event = (wifi_event_ap_stadisconnected_t*) event_data;
    ESP_LOGI(TAG, "station "MACSTR" leave, AID=%d", MAC2STR(event->mac), event->aid);
  }
}

static esp_err_t wifi_config_get(const char *path, sys_wifi_cfg_t *cfg)
{
  int rd_bytes = 0;
  esp_err_t err = ESP_OK;
  memset(cfg->host_name, 0, 20);
  memset(cfg->ap_ssid, 0, WIFI_SSID_PASSWD_MAX_LEN);
  memset(cfg->ap_passwd, 0, WIFI_SSID_PASSWD_MAX_LEN);
  memset(cfg->sta_ssid, 0, WIFI_SSID_PASSWD_MAX_LEN);
  memset(cfg->sta_passwd, 0, WIFI_SSID_PASSWD_MAX_LEN);
  strcpy(cfg->host_name, "myled");
  strcpy(cfg->ap_ssid, "kychu_light");
  strcpy(cfg->ap_passwd, "1234567890");

  char filepath[32];
  memset(filepath, 0, 32);
  strlcpy(filepath, path, 32);
  strlcat(filepath, "/wifi.json", 32);

  FILE *fp = fopen(filepath, "r");
  if(fp == NULL) {
	ESP_LOGW(TAG, "open %s failed, create it!", filepath);
    fp = fopen(filepath, "w");
    if(fp == NULL) {
      ESP_LOGE(TAG, "%s create failed!", filepath);
      err = ESP_ERR_NOT_FOUND;
    } else {
      cJSON *root = cJSON_CreateObject();
      cJSON_AddStringToObject(root, wifi_cfg_json_name[0], cfg->host_name);
      cJSON_AddStringToObject(root, wifi_cfg_json_name[1], cfg->sta_ssid);
      cJSON_AddStringToObject(root, wifi_cfg_json_name[2], cfg->sta_passwd);
      cJSON_AddStringToObject(root, wifi_cfg_json_name[3], cfg->ap_ssid);
      cJSON_AddStringToObject(root, wifi_cfg_json_name[4], cfg->ap_passwd);
      const char *sys_info = cJSON_Print(root);
      int len = strlen(sys_info);
      if((len > 0) && len != fwrite(sys_info, 1, len, fp)) {
        ESP_LOGE(TAG, "failed to save configuration");
        unlink(filepath);
        len = -1;
      }
      free((void *)sys_info);
      cJSON_Delete(root);
      fclose(fp);
      if(len < 0) err = ESP_FAIL;
    }
  } else {
	char *buff = calloc(1, 1024);
	if(buff != NULL) {
	  if((rd_bytes = fread(buff, 1, 1024, fp)) == -1) {
		ESP_LOGE(TAG, "failed to read %s", filepath);
		err = ESP_FAIL;
	  } else {
		buff[rd_bytes] = 0;
		cJSON *root = cJSON_Parse(buff);
		strcpy(cfg->host_name, cJSON_GetObjectItem(root, wifi_cfg_json_name[0])->valuestring);
		ESP_LOGI(TAG, "name: %s", cfg->host_name);
		strcpy(cfg->sta_ssid, cJSON_GetObjectItem(root, wifi_cfg_json_name[1])->valuestring);
		ESP_LOGI(TAG, "sta_ssid: %s", cfg->sta_ssid);
		strcpy(cfg->sta_passwd, cJSON_GetObjectItem(root, wifi_cfg_json_name[2])->valuestring);
		ESP_LOGI(TAG, "sta_passwd: %s", cfg->sta_passwd);
		strcpy(cfg->ap_ssid, cJSON_GetObjectItem(root, wifi_cfg_json_name[3])->valuestring);
		ESP_LOGI(TAG, "ap_ssid: %s", cfg->ap_ssid);
		strcpy(cfg->ap_passwd, cJSON_GetObjectItem(root, wifi_cfg_json_name[4])->valuestring);
		ESP_LOGI(TAG, "ap_passwd: %s", cfg->ap_passwd);
		cJSON_Delete(root);
	  }
	  free(buff);
	} else {
	  ESP_LOGE(TAG, "no memory for file read cache");
	  err = ESP_ERR_NO_MEM;
	}
	fclose(fp);
  }
  return err;
}

#define MDNS_INSTANCE "led light web server"
static void initialise_mdns(const char *host_name)
{
  mdns_init();
  mdns_hostname_set(host_name);
  mdns_instance_name_set(MDNS_INSTANCE);

  ESP_ERROR_CHECK(mdns_service_add("Light-WebServer", "_http", "_tcp", 80, NULL, 0));
}

esp_err_t wifi_connect_init(const char *path)
{
  cfg = calloc(1, sizeof(sys_wifi_cfg_t));
  if(cfg == NULL) return ESP_ERR_NO_MEM;
  int sta_initialized = 0;

  ESP_ERROR_CHECK(wifi_config_get(path, cfg));

  ESP_ERROR_CHECK(esp_netif_init());
  ESP_ERROR_CHECK(esp_event_loop_create_default());
  esp_netif_create_default_wifi_sta();
  esp_netif_create_default_wifi_ap();

  wifi_init_config_t wcfg = WIFI_INIT_CONFIG_DEFAULT();
  ESP_ERROR_CHECK(esp_wifi_init(&wcfg));

  wifi_config_t wifi_config;

  ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
                                                      ESP_EVENT_ANY_ID,
                                                      &event_handler,
                                                      NULL, NULL));

  if(strlen(cfg->sta_ssid) != 0 && strlen(cfg->sta_passwd) != 0) { // sta_ssid are available
    s_wifi_event_group = xEventGroupCreate();

    esp_event_handler_instance_t instance_got_ip;
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
                                                        IP_EVENT_STA_GOT_IP,
                                                        &event_handler,
                                                        NULL,
                                                        &instance_got_ip));


    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    memset(&wifi_config, 0, sizeof(wifi_config_t));
    /* Setting a password implies station will connect to all security modes including WEP/WPA.
     * However these modes are deprecated and not advisable to be used. Incase your Access point
     * doesn't support WPA2, these mode can be enabled by commenting below line */
    wifi_config.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
    wifi_config.sta.pmf_cfg.capable = true;
    wifi_config.sta.pmf_cfg.required = false;
    wifi_config.sta.bssid_set = false;
    strcpy((char *)wifi_config.sta.ssid, cfg->sta_ssid);
    strcpy((char *)wifi_config.sta.password, cfg->sta_passwd);
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    /* Waiting until either the connection is established (WIFI_CONNECTED_BIT) or connection failed for the maximum
     * number of re-tries (WIFI_FAIL_BIT). The bits are set by event_handler() (see above) */
    EventBits_t bits = xEventGroupWaitBits(s_wifi_event_group,
                                           WIFI_CONNECTED_BIT | WIFI_FAIL_BIT,
                                           pdFALSE,
                                           pdFALSE,
                                           portMAX_DELAY);

    /* xEventGroupWaitBits() returns the bits before the call returned, hence we can test which event actually
     * happened. */
    if (bits & WIFI_CONNECTED_BIT) {
      ESP_LOGI(TAG, "connected to ap SSID:%s password:%s", cfg->sta_ssid, cfg->sta_passwd);
      ESP_LOGI(TAG, "wifi station initialized.");
      sta_initialized = 1;
    } else if (bits & WIFI_FAIL_BIT) {
      ESP_LOGW(TAG, "Failed to connect to SSID:%s, password:%s", cfg->sta_ssid, cfg->sta_passwd);
      ESP_LOGW(TAG, "start AP mode");
    } else {
      ESP_LOGE(TAG, "UNEXPECTED EVENT");
    }

    /* The event will not be processed after unregister */
    ESP_ERROR_CHECK(esp_event_handler_instance_unregister(IP_EVENT, IP_EVENT_STA_GOT_IP, instance_got_ip));
    vEventGroupDelete(s_wifi_event_group);

    if(sta_initialized) {
      initialise_mdns(cfg->host_name);
      netbiosns_init();
      netbiosns_set_name(cfg->host_name);
    }
  }
  if(!sta_initialized) {
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));
    memset(&wifi_config, 0, sizeof(wifi_config_t));
    strlcpy((char *)wifi_config.ap.ssid, cfg->ap_ssid, sizeof(wifi_config.ap.ssid));
    wifi_config.ap.ssid_len = strlen((const char *)wifi_config.ap.ssid);
    wifi_config.ap.channel = EXAMPLE_ESP_WIFI_CHANNEL;
    if(strlen(cfg->ap_passwd) == 0) {
      wifi_config.ap.authmode = WIFI_AUTH_OPEN;
    } else {
      strlcpy((char *)wifi_config.ap.password, cfg->ap_passwd, sizeof(wifi_config.ap.password));
      wifi_config.ap.authmode = WIFI_AUTH_WPA_WPA2_PSK;
    }
    wifi_config.ap.max_connection = EXAMPLE_MAX_STA_CONN;
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());
    ESP_LOGI(TAG, "wifi softAP initialized. (%s:%s)", wifi_config.ap.ssid, wifi_config.ap.password);
  }

  return ESP_OK;
}

const char *wifi_get_hostname(void)
{
  if(cfg == NULL) return NULL;
  return (const char *)cfg->host_name;
}

const char *wifi_get_ap_ssid(void)
{
  if(cfg == NULL) return NULL;
  return (const char *)cfg->ap_ssid;
}

const char *wifi_get_ap_passwd(void)
{
  if(cfg == NULL) return NULL;
  return (const char *)cfg->ap_passwd;
}

const char *wifi_get_sta_ssid(void)
{
  if(cfg == NULL) return NULL;
  return (const char *)cfg->sta_ssid;
}

const char *wifi_get_sta_passwd(void)
{
  if(cfg == NULL) return NULL;
  return (const char *)cfg->sta_passwd;
}

const char *wifi_get_sta_ip(void)
{
  return (const char *)ipv4_ip_str;
}

esp_err_t wifi_config_validate(const char *path)
{
  char filepath[32];
  memset(filepath, 0, 32);
  strlcpy(filepath, path, 32);
  strlcat(filepath, "/wifi.json", 32);

  FILE *fp = fopen(filepath, "w");
  if(fp == NULL) {
    ESP_LOGE(TAG, "failed to open %s", filepath);
    return ESP_ERR_NOT_FOUND;
  } else {
    cJSON *root = cJSON_CreateObject();
    cJSON_AddStringToObject(root, wifi_cfg_json_name[0], cfg->host_name);
    cJSON_AddStringToObject(root, wifi_cfg_json_name[1], cfg->sta_ssid);
    cJSON_AddStringToObject(root, wifi_cfg_json_name[2], cfg->sta_passwd);
    cJSON_AddStringToObject(root, wifi_cfg_json_name[3], cfg->ap_ssid);
    cJSON_AddStringToObject(root, wifi_cfg_json_name[4], cfg->ap_passwd);
    const char *sys_info = cJSON_Print(root);
    int len = strlen(sys_info);
    if((len > 0) && len != fwrite(sys_info, 1, len, fp)) {
      ESP_LOGE(TAG, "failed to save configuration");
      unlink(filepath);
      len = -1;
    }
    free((void *)sys_info);
    cJSON_Delete(root);
    fclose(fp);
    if(len < 0) return ESP_FAIL;
    return ESP_OK;
  }
}

esp_err_t wifi_set_hostname(const char *s)
{
  if(cfg == NULL || s == NULL || strlen(s) < 6) return ESP_FAIL;
  memset(cfg->host_name, 0, WIFI_SSID_PASSWD_MAX_LEN);
  strlcpy(cfg->host_name, s, WIFI_SSID_PASSWD_MAX_LEN);
  return ESP_OK;
}

esp_err_t wifi_set_ap_ssid(const char *s)
{
  if(cfg == NULL || s == NULL || strlen(s) < 6) return ESP_FAIL;
  memset(cfg->ap_ssid, 0, WIFI_SSID_PASSWD_MAX_LEN);
  strlcpy(cfg->ap_ssid, s, WIFI_SSID_PASSWD_MAX_LEN);
  return ESP_OK;
}

esp_err_t wifi_set_ap_passwd(const char *s)
{
  if(cfg == NULL || s == NULL || strlen(s) < 8) return ESP_FAIL;
  memset(cfg->ap_passwd, 0, WIFI_SSID_PASSWD_MAX_LEN);
  strlcpy(cfg->ap_passwd, s, WIFI_SSID_PASSWD_MAX_LEN);
  return ESP_OK;
}

esp_err_t wifi_set_sta_ssid(const char *s)
{
  if(cfg == NULL || s == NULL || strlen(s) < 6) return ESP_FAIL;
  memset(cfg->sta_ssid, 0, WIFI_SSID_PASSWD_MAX_LEN);
  strlcpy(cfg->sta_ssid, s, WIFI_SSID_PASSWD_MAX_LEN);
  return ESP_OK;
}

esp_err_t wifi_set_sta_passwd(const char *s)
{
  if(cfg == NULL || s == NULL || strlen(s) < 8) return ESP_FAIL;
  memset(cfg->sta_passwd, 0, WIFI_SSID_PASSWD_MAX_LEN);
  strlcpy(cfg->sta_passwd, s, WIFI_SSID_PASSWD_MAX_LEN);
  return ESP_OK;
}
